use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;

use dashmap::DashMap;
use tokio::net::TcpListener;
use tokio::sync::{mpsc, RwLock};
use tokio::time::timeout;
use tonic::{Request, Response, Status};
use tonic::transport::Server;
use tracing::{debug, error, info, warn};
use uuid::Uuid;
use tokio_stream::wrappers::ReceiverStream;

use crate::actor::{Actor, ActorError, PID, Props};
use crate::protos::remote::{
    MessageBatch, MessageEnvelope, RemoteMessage, 
    ActorPidRequest, ActorPidResponse,
    ListProcessesRequest, ListProcessesResponse,
    GetProcessDiagnosticsRequest, GetProcessDiagnosticsResponse,
};
use crate::protos::remote::remoting_server::{Remoting, RemotingServer as TonicRemotingServer};
use crate::remote::{RemoteError, pid_to_proto, proto_to_pid};
use crate::remote::config::RemoteConfig;
use crate::remote::endpoint_manager::EndpointManager;
use crate::remote::serialization::{SerializerRegistry, ProtoSerializer};
use crate::system::ActorSystem;

/// Remote actor system server
pub struct RemoteServer {
    /// Reference to the actor system
    system: ActorSystem,
    
    /// Remote configuration
    config: Arc<RemoteConfig>,
    
    /// Endpoint manager
    endpoint_manager: Arc<EndpointManager>,
    
    /// Serializer registry
    serializer_registry: Arc<SerializerRegistry>,
    
    /// Map of actor kind names to actor props
    kinds: Arc<DashMap<String, Props>>,
    
    /// System ID
    system_id: String,
    
    /// Server shutdown channel
    shutdown_tx: Option<mpsc::Sender<()>>,
}

impl RemoteServer {
    /// Create a new remote server
    pub fn new(system: ActorSystem, config: RemoteConfig) -> Self {
        let config = Arc::new(config);
        
        // Initialize the kinds map with the registered actor types
        let kinds = Arc::new(DashMap::new());
        for (kind, props) in config.kinds.as_ref().iter() {
            kinds.insert(kind.clone(), props.clone());
        }
        
        // Generate a unique system ID
        let system_id = Uuid::new_v4().to_string();
        
        Self {
            system: system.clone(),
            config: config.clone(),
            endpoint_manager: Arc::new(EndpointManager::new(system, (*config.clone()).clone())),
            serializer_registry: Arc::new(SerializerRegistry::new()),
            kinds,
            system_id,
            shutdown_tx: None,
        }
    }
    
    /// Start the remote server
    pub async fn start(&mut self) -> Result<(), RemoteError> {
        info!("Starting remote server on {}:{}", self.config.host, self.config.port);
        
        // Create a shutdown channel
        let (tx, mut rx) = mpsc::channel::<()>(1);
        self.shutdown_tx = Some(tx);
        
        // Clone references for the server task
        let addr = format!("{}:{}", self.config.host, self.config.port);
        let endpoint_manager = self.endpoint_manager.clone();
        let serializer_registry = self.serializer_registry.clone();
        let kinds = self.kinds.clone();
        let system = self.system.clone();
        let system_id = self.system_id.clone();
        
        // Start the gRPC server in a separate task
        tokio::spawn(async move {
            // Create the service implementation
            let remote_service = RemoteServiceImpl {
                endpoint_manager: endpoint_manager.clone(),
                serializer_registry: serializer_registry.clone(),
                kinds: kinds.clone(),
                system: system.clone(),
                system_id: system_id.clone(),
            };
            
            // Create the server
            let server = Server::builder()
                .add_service(TonicRemotingServer::new(remote_service))
                .serve_with_shutdown(
                    addr.parse().unwrap(),
                    async {
                        rx.recv().await;
                        info!("Remote server shutdown signal received");
                    }
                );
            
            // Start the server
            info!("Remote server listening on {}", addr);
            if let Err(err) = server.await {
                error!("Remote server error: {}", err);
            }
        });
        
        Ok(())
    }
    
    /// Stop the remote server
    pub async fn stop(&mut self) -> Result<(), RemoteError> {
        info!("Stopping remote server");
        
        // Stop all endpoints
        self.endpoint_manager.stop_all();
        
        // Send shutdown signal to the server
        if let Some(tx) = self.shutdown_tx.take() {
            let _ = tx.send(()).await;
        }
        
        info!("Remote server stopped");
        Ok(())
    }
    
    /// Register an actor kind that can be spawned remotely
    pub fn register<A: Actor + 'static>(&self, kind: &str, props: Props) -> Result<(), RemoteError> {
        info!("Registering actor kind: {}", kind);
        self.kinds.insert(kind.to_string(), props);
        Ok(())
    }
    
    /// Spawn an actor remotely
    pub async fn spawn_remote<A: Actor + 'static>(
        &self,
        address: &str,
        kind: &str,
        name: Option<&str>,
        timeout_duration: Duration,
    ) -> Result<PID, RemoteError> {
        info!("Spawning remote actor: kind={}, address={}", kind, address);
        
        // Get or create an endpoint for the address
        let endpoint = self.endpoint_manager.get_or_create_endpoint(address);
        
        // Create a PID request
        let request = ActorPidRequest {
            kind: kind.to_string(),
            name: name.unwrap_or("").to_string(),
        };
        
        // In a real implementation, we'd send the request to the remote endpoint
        // and wait for a response with the PID
        
        // For now, we'll just create a PID with a random ID
        let id = Uuid::new_v4().to_string();
        let pid = PID {
            address: address.to_string(),
            id,
        };
        
        Ok(pid)
    }
}

/// Implementation of the remote actor service
pub struct RemoteServiceImpl {
    /// Endpoint manager
    endpoint_manager: Arc<EndpointManager>,
    
    /// Serializer registry
    serializer_registry: Arc<SerializerRegistry>,
    
    /// Map of actor kind names to actor props
    kinds: Arc<DashMap<String, Props>>,
    
    /// Reference to the actor system
    system: ActorSystem,
    
    /// System ID
    system_id: String,
}

#[tonic::async_trait]
impl Remoting for RemoteServiceImpl {
    /// Bidirectional streaming for message exchange
    type ReceiveStream = ReceiverStream<Result<RemoteMessage, Status>>;
    
    async fn receive(
        &self,
        request: Request<tonic::Streaming<RemoteMessage>>,
    ) -> Result<Response<Self::ReceiveStream>, Status> {
        info!("New remote connection established");
        
        // Get the client's stream
        let mut stream = request.into_inner();
        
        // Create a channel for our responses
        let (tx, rx) = mpsc::channel(100);
        
        // Clone references for the receiver task
        let endpoint_manager = self.endpoint_manager.clone();
        let serializer_registry = self.serializer_registry.clone();
        let system = self.system.clone();
        
        // Spawn a task to handle incoming messages
        tokio::spawn(async move {
            while let Some(message) = stream.message().await.transpose() {
                match message {
                    Ok(msg) => {
                        // Process the message
                        debug!("Received remote message");
                        
                        // In a real implementation, we'd create an EndpointReader
                        // and process the message batch
                    }
                    Err(err) => {
                        error!("Error receiving message: {}", err);
                        break;
                    }
                }
            }
            
            info!("Remote connection closed");
        });
        
        Ok(Response::new(ReceiverStream::new(rx)))
    }
    
    /// List processes matching a pattern
    async fn list_processes(
        &self,
        request: Request<ListProcessesRequest>,
    ) -> Result<Response<ListProcessesResponse>, Status> {
        let req = request.into_inner();
        info!("List processes request: pattern={}", req.pattern);
        
        // In a real implementation, we'd query the process registry
        // and filter the results based on the pattern
        
        Ok(Response::new(ListProcessesResponse {
            pids: Vec::new(),
        }))
    }
    
    /// Get diagnostics info for a process
    async fn get_process_diagnostics(
        &self,
        request: Request<GetProcessDiagnosticsRequest>,
    ) -> Result<Response<GetProcessDiagnosticsResponse>, Status> {
        let req = request.into_inner();
        let pid = req.pid.ok_or_else(|| Status::invalid_argument("PID is required"))?;
        let pid = proto_to_pid(&pid);
        
        info!("Get process diagnostics request: pid={}", pid);
        
        // In a real implementation, we'd query the process for diagnostics
        
        Ok(Response::new(GetProcessDiagnosticsResponse {
            diagnostics_string: format!("Diagnostics for {}", pid),
        }))
    }
}

/// Main remote actor system implementation
pub struct Remote {
    /// The remote server
    server: RemoteServer,
    
    /// Whether the system is started
    started: bool,
}

impl Remote {
    /// Create a new remote actor system
    pub fn new(system: ActorSystem, config: RemoteConfig) -> Self {
        Self {
            server: RemoteServer::new(system, config),
            started: false,
        }
    }
    
    /// Start the remote actor system
    pub async fn start(&mut self) -> Result<(), RemoteError> {
        if self.started {
            return Ok(());
        }
        
        self.server.start().await?;
        self.started = true;
        
        Ok(())
    }
    
    /// Stop the remote actor system
    pub async fn stop(&mut self) -> Result<(), RemoteError> {
        if !self.started {
            return Ok(());
        }
        
        self.server.stop().await?;
        self.started = false;
        
        Ok(())
    }
    
    /// Register an actor kind that can be spawned remotely
    pub fn register<A: Actor + 'static>(&self, kind: &str, props: Props) -> Result<(), RemoteError> {
        self.server.register::<A>(kind, props)
    }
    
    /// Spawn an actor remotely
    pub async fn spawn_remote<A: Actor + 'static>(
        &self,
        address: &str,
        kind: &str,
        name: Option<&str>,
        timeout_duration: Duration,
    ) -> Result<PID, RemoteError> {
        self.server.spawn_remote::<A>(address, kind, name, timeout_duration).await
    }
} 